import torch.nn as nn
import torch
import torch.nn.functional as F

class HeatmapLoss(nn.Module):
    def __init__(self,  alpha=2, beta=4, reduction='mean'):
        super(HeatmapLoss, self).__init__()
        self.alpha = alpha
        self.beta = beta
    def forward(self, inputs, targets):
        inputs = F.sigmoid(inputs)
        center_id = (targets >= 0.95).float()
        other_id = (targets < 0.95).float()
        center_loss = -center_id * (1.0-inputs)**self.alpha * torch.log(inputs + 1e-14)
        other_loss = -other_id * (1 - targets)**self.beta * (inputs)**self.alpha * torch.log(1.0 - inputs + 1e-14)

        return torch.sum(center_loss + other_loss)

